#!/usr/bin/env python3
"""
run.py - master runner for Political DEBATE replication archive.

This script:
 - Accepts command-line options to choose hardware (cpu, mps, cuda) and which sections to run
   (analysis, labeling, gpu_benchmarks, cpu_benchmarks, mps_benchmarks, llama).
 - Finds `.py` and `.R` scripts inside section directories (alphabetical order).
 - Runs each script, streaming both stdout and stderr to the console and to a timestamped log file.
 - Exposes the selected hardware via the DEVICE environment variable for child scripts to read.
 - Exits with the return code of the first failing script (unless --continue-on-error is set).
"""

import argparse
import subprocess
import sys
import os
from pathlib import Path
from datetime import datetime
import shutil

# Root is the directory containing this file. We assume this is the archive root.
ROOT = Path(__file__).resolve().parent

# Allowed section names and hardware choices
VALID_SECTIONS = ["analysis", "labeling", "gpu_benchmarks", "cpu_benchmarks", "mps_benchmarks", "llama"]
VALID_HARDWARE = ["cpu", "mps", "cuda"]

def make_log_file():
    """
    Create a logs directory (if needed) and return a Path for a new timestamped log file.
    This ensures every run writes to a unique logfile.
    """
    logs_dir = ROOT / "logs"
    logs_dir.mkdir(exist_ok=True)
    ts = datetime.now().strftime("%Y%m%d_%H%M%S")
    return logs_dir / f"run_{ts}.log"

def check_prereqs():
    """
    Light prereq check: return the path to 'Rscript' if available on PATH, otherwise None.
    We use this before running sections to warn the user if R is required but not installed.
    """
    return shutil.which("Rscript")

def find_scripts_for_section(section_dir: Path):
    """
    Return a list of scripts in a section directory sorted alphabetically.
    We look for Python (.py) files and R (.R/.r) files.
    The return value is a list of tuples (path, type) where type is 'py' or 'r'.

    This alphabetical ordering gives you deterministic ordering and lets you control order
    by prefixing files (e.g., 01_data_prep.py, 02_analysis.R).
    """
    if not section_dir.exists():
        return []
    files = []
    for ext in ("*.py", "*.R", "*.r"):
        files.extend(sorted(section_dir.glob(ext)))
    normalized = []
    for f in files:
        if f.suffix.lower() == ".py":
            normalized.append((f, "py"))
        else:
            normalized.append((f, "r"))
    return normalized

def stream_subprocess(cmd, env, cwd, logfile, label):
    """
    Run a subprocess (cmd list) and stream its stdout/stderr both to the console and to `logfile`.

    - cmd: list (executable + args)
    - env: environment dict for the subprocess
    - cwd: working directory for the subprocess
    - logfile: open file object where we append the output
    - label: textual label used in the logs to identify this command

    Returns the subprocess return code.
    """
    # Write header info for this run into the log
    logfile.write(f"\n\n>>> Running: {label}\n>>> Command: {' '.join(cmd)}\n")
    logfile.flush()

    # Start the process, capturing combined stdout+stderr so we can stream them in one loop.
    proc = subprocess.Popen(
        cmd,
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
        text=True,
        env=env,
        cwd=cwd
    )

    # Iterate over the subprocess stdout stream line-by-line.
    # This allows live streaming to console and appending to log.
    for line in proc.stdout:
        print(line, end="")          # stream to console
        logfile.write(line)          # append to logfile
        logfile.flush()

    proc.wait()  # wait for process to finish
    logfile.write(f"\n>>> Exit code: {proc.returncode}\n")
    logfile.flush()
    return proc.returncode

def run_section(section, hardware, logfile, continue_on_error=False):
    """
    Run every script inside a named section folder.

    - section: one of VALID_SECTIONS (e.g., "analysis")
    - hardware: string like "cpu" / "gpu" / "mps" (used to set DEVICE in environment)
    - logfile: open file to append logs
    - continue_on_error: if True, don't abort the whole section on the first failing script

    Behavior:
      - For .py files: runs them with the same Python interpreter running run.py.
      - For .R files: runs them with Rscript (looks it up with shutil.which).
      - If a script returns non-zero, this function returns that code (unless continue_on_error True,
        in which case it continues and returns 0 if all scripts succeed, or the last non-zero code).
    """
    section_dir = ROOT / section
    scripts = find_scripts_for_section(section_dir)
    if not scripts:
        logfile.write(f"\n-- No scripts found for section: {section} (looked in {section_dir})\n")
        logfile.flush()
        return 0

    logfile.write(f"\n-- Running section: {section} (scripts found: {len(scripts)})\n")
    logfile.flush()

    # Prepare environment for child scripts: copy current environment and inject DEVICE
    base_env = os.environ.copy()
    base_env["DEVICE"] = hardware

    # If user explicitly requested CPU, hide CUDA_VISIBLE_DEVICES so child scripts don't find GPUs
    if hardware == "cpu":
        base_env["CUDA_VISIBLE_DEVICES"] = ""

    # Iterate through found scripts in order
    for script_path, typ in scripts:
        label = f"{section}/{script_path.name}"

        # Choose the command to run based on file type
        if typ == "py":
            # Use the same Python executable that is running run.py (ensures consistent env)
            cmd = [sys.executable, str(script_path)]
        else:
            # For R scripts, find Rscript. If missing, log an error and return non-zero.
            rscript_path = shutil.which("Rscript")
            if not rscript_path:
                logfile.write("ERROR: Rscript not found on PATH but an R script needs to run.\n")
                logfile.flush()
                return 2
            cmd = [rscript_path, str(script_path)]

        # Log start time and run the script, streaming output
        logfile.write(f"\n--- Starting {label} at {datetime.now().isoformat()} ---\n")
        logfile.flush()
        rc = stream_subprocess(cmd=cmd, env=base_env, cwd=ROOT, logfile=logfile, label=label)
        logfile.write(f"--- Finished {label} at {datetime.now().isoformat()} (rc={rc}) ---\n")
        logfile.flush()

        # Handle non-zero exit codes according to continue_on_error
        if rc != 0:
            errmsg = f"Script {label} failed with return code {rc}."
            logfile.write(errmsg + "\n")
            logfile.flush()
            print(errmsg, file=sys.stderr)
            if not continue_on_error:
                # Return immediately so main() can abort
                return rc
            # if continue_on_error is True we simply continue to next script
    return 0

def main():
    """
    Parse command-line arguments, prepare logging, optionally show what will run (list-only),
    then run requested sections in order. Exits with the final return code (0 on success).
    """
    parser = argparse.ArgumentParser(description="Run replication archive sections (python + R).")
    parser.add_argument("--hardware", choices=VALID_HARDWARE, default="cpu",
                        help="hardware target to pass to scripts (DEVICE env var).")
    parser.add_argument("--sections", nargs="+", choices=VALID_SECTIONS, default=["analysis"],
                        help="which sections to run (space-separated).")
    parser.add_argument("--continue-on-error", action="store_true",
                        help="if set, continue running remaining scripts after an error.")
    parser.add_argument("--list-only", action="store_true",
                        help="only list scripts that would be run (no execution).")
    args = parser.parse_args()

    # Create a new log file for this run
    logfile_path = make_log_file()
    with open(logfile_path, "w", encoding="utf-8") as logfile:
        # Write some run metadata to the log
        logfile.write(f"Replication run started: {datetime.now().isoformat()}\n")
        logfile.write(f"Root: {ROOT}\n")
        logfile.write(f"Hardware: {args.hardware}\n")
        logfile.write(f"Sections: {args.sections}\n")
        logfile.write(f"Continue on error: {args.continue_on_error}\n")
        logfile.flush()

        # If the user only asked to list scripts, do that and exit early
        if args.list_only:
            for sec in args.sections:
                sec_dir = ROOT / sec
                scripts = find_scripts_for_section(sec_dir)
                logfile.write(f"\n-- Section: {sec} ({len(scripts)} scripts)\n")
                print(f"Section: {sec}")
                for p, t in scripts:
                    print(f"  - {sec}/{p.name}")
                    logfile.write(f"  - {sec}/{p.name}\n")
            logfile.write("\nList-only mode, exiting.\n")
            logfile.flush()
            print(f"\nLog saved to {logfile_path}")
            return

        # Pre-check: if any selected sections contain R scripts but Rscript is missing, warn the user.
        # This is a convenience check; run_section will also check before running each R script.
        need_r = False
        for sec in args.sections:
            scripts = find_scripts_for_section(ROOT / sec)
            if any(t == "r" for (_, t) in scripts):
                need_r = True
                break
        if need_r and not check_prereqs():
            logfile.write("WARNING: Rscript not found on PATH but some sections contain R scripts.\n")
            logfile.flush()
            print("WARNING: Rscript not found on PATH but some sections contain R scripts.", file=sys.stderr)

        # Run requested sections in the order they were provided on the command line
        final_rc = 0
        for sec in args.sections:
            rc = run_section(sec, args.hardware, logfile, continue_on_error=args.continue_on_error)
            if rc != 0:
                final_rc = rc
                if not args.continue_on_error:
                    logfile.write(f"Aborting due to failure in section {sec} (rc={rc}).\n")
                    logfile.flush()
                    break

        logfile.write(f"\nReplication run finished: {datetime.now().isoformat()} (final_rc={final_rc})\n")
        logfile.flush()

    # Print the path to the logfile for the user's convenience and exit with the result code
    print(f"Run complete. Log file: {logfile_path}")
    sys.exit(final_rc)

if __name__ == "__main__":
    main()
